
from chains.C0_docx_limit import get_word_page_limit
from chains.C1_word2json import word2json
from chains.C24_overall_bubble_v2 import Overall_Bubble_V2
from chains.C2_recall_material import recall_material_list
from chains.C3_material_problem import Material2Problem
from chains.C4_get_prompt import get_prompt
from chains.C5_AI_answer import AI_Answer
from chains.C6_detail_by_problrm import detail_by_problem
from chains.C7_material_full_text import material_full_text
from chains.C8_convert_word import split_word_to_run
from chains.C13_bubble_score import Global_Score
from chains.C15_bubble_title import Bubble
from chains.C20_get_prompt_third import get_prompt_third
from chains.C24_overall_bubble import Overall_Bubble
from chains.C26_slice_query_title import slice_query_title
from chains.C27_text_correction import Text_Correction
from config.config import config
from chains import C4_get_prompt_v2
from pdf2paragraph.pdf_main import get_full_text
from fastapi.responses import StreamingResponse
from fastapi import APIRouter, HTTPException, Request
from fastapi import FastAPI, File, UploadFile, Body,BackgroundTasks
from fastapi.responses import FileResponse
from fastapi.responses import JSONResponse
from fastapi.responses import Response
from pydantic import BaseModel, Field
from collections import OrderedDict
from typing import Optional

from utils.IPOBaseResponse import IPOBaseResponse
from utils.db import LawDB
from docx import Document
from pathlib import Path
from typing import List
from io import BytesIO
import requests
import shutil
import json
import time
import io
import os
import re
import chains.C4_get_prompt_v2
from utils import db_pool


router = APIRouter(tags=["IPO_Review_API"], responses={404: {"description": "Not Found"}})

class Data_C0(BaseModel):
    # file_path: str = Field(default="./sentence_docx/fileId.docx")
    fileId: str = Field(default="")
    DownloadDataName: str = Field(default="")
    l_sign: str = Field(default="")
    Authorization: str = Field(default="")


@router.post("/C0_Docx_Limit", summary="C0：输入: word参数 => 输出: true 或者 false", include_in_schema=True)
async def C0_docx_limit(data: Data_C0):
    page_limit = await get_word_page_limit(data.fileId, data.DownloadDataName, data.l_sign, data.Authorization)
    return page_limit


class Data_C1(BaseModel):
    file_path: str = Field(default="./sentence_docx/fileId.docx", description="可渲染的docx路径",)

@router.post("/C1_Word2Json", summary="C1：输入: word => 输出: 切片后的json", include_in_schema=True)
async def C1_word2json_request(data: Data_C1):
    index_paragraph_sentence_dict = await word2json(data.file_path)
    return index_paragraph_sentence_dict

class Data_C2(BaseModel):
    content_text: str = Field(default="")
    k: int = Field(default=3)
    sector: str = Field(default="板块")
    industry: str = Field(default="行业")
    region: Optional[str] = Field(default="地域")
    problem_type_list: List[str] = Field(default_factory=list)


@router.post("/C2_Recall_Material", summary="C2：输入: 切片、k、板块、行业、地域 => 输出: [[{材料ID: 材料内容}]]", include_in_schema=True)
async def C2_recall_material(data: Data_C2):

    material_list = await recall_material_list(data.content_text, data.k, data.sector, data.industry, data.region, data.problem_type_list)

    return material_list


class Data_C3(BaseModel):
    material_id_list: List[int] = Field(default_factory=list)

@router.post("/C3_Material_Problem",summary="C3：输入: 一维的材料ID列表 => 输出: 一维的问题ID列表", include_in_schema=True)
async def C3_material_problem_request(data: Data_C3):
    Mater = Material2Problem()
    problem_id_list = await Mater.material_problem(data.material_id_list)
    return problem_id_list


class Data_C4(BaseModel):
    content_text: str = Field(default="")
    problem_id_list: List[int] = Field(default_factory=list)
    simple: bool = Field(default=False)

@router.post("/C4_Get_Prompt", summary="C4：输入: 招股书文本 和 问题ID列表 => 输出: 提示词", include_in_schema=True)
async def C4_get_prompt_request(data: Data_C4):
    prompt = await get_prompt(data.content_text, data.problem_id_list)
    return prompt

@router.post("/C4_Get_Prompt_V2", summary="C4_V2：输入: 招股书文本 和 问题ID列表 => 输出: 提示词", include_in_schema=True)
async def C4_get_prompt_request(data: Data_C4):
    """
    选中某段文字后, 得到问询，然后输入到大模型
    :param data:
    :return:
    """
    category_path = os.path.join(config.DATA_DIR, 'rules_sample.json')
    rules_path = os.path.join(config.DATA_DIR, 'rules_query_merged.json')
    ai_client = AI_Answer()
    category_list = await C4_get_prompt_v2.get_material_category(ai_client, data.content_text, category_path)
    prompt = await C4_get_prompt_v2.get_prompt(ai_client, data.content_text, category_list, rules_path, data.simple)
    return prompt


class Data_C5(BaseModel):
    prompt: str = Field(default="")

@router.post("/C5_AI_Answer", summary="C5：输入: 提示词 => 输出: AI回答", include_in_schema=True)
async def C5_AI_answer_request(data: Data_C5):
    AI = AI_Answer()
    return StreamingResponse(AI.get_ai_answer_R1(data.prompt), media_type="text/event-stream")


class Data_C6(BaseModel):
    problem_id_list: List[int] = Field(default_factory=list)
    sorting_type: str = Field(default="相似度降序", description="日期降序, 日期升序, 相似度降序")

@router.post("/C6_Detail_By_Problem", summary="C6：输入: 问题ID列表 => 输出:详情列表", include_in_schema=True)
async def C6_detail_by_problem_request(data: Data_C6):
    detail_list = await detail_by_problem(data.problem_id_list, data.sorting_type)
    return detail_list


class Data_C7(BaseModel):
    material_id: int = Field(default=1)

@router.post("/C7_Material_Full_Text", summary="C7：输入: 材料ID => 输出:该pdf所有问题ID列表", include_in_schema=True)
async def C7_material_full_text_request(data: Data_C7):
    full_text_problem_id_list = await material_full_text(data.material_id)
    return full_text_problem_id_list


class Data_C8(BaseModel):
    fileId: str = Field(default="文件加密id")
    DownloadDataName: str = Field(default="文件加密参数")
    l_sign: str = Field(default="文件加密参数")
    Authorization: str = Field(default="文件加密参数")

@router.post("/C8_Convert_Word", summary="C8：转换docx为按句子处理后的docx", include_in_schema=True)
async def C8_convert_word(Data: Data_C8):
    s_time = time.time()
    # 如果句子docx已经存在 则直接返回文件流
    file_path = Path("./sentence_docx") / f"{Data.fileId}.docx"
    if os.path.exists(file_path):
        document = Document(file_path)
        new_file_stream = io.BytesIO()
        document.save(new_file_stream)
        new_file_stream.seek(0)
        # 使用 getvalue() 方法获取字节流
        bytes_stream_again = new_file_stream.getvalue()
        return Response(content=bytes_stream_again, media_type="application/octet-stream")

    # 接口 URL
    url = 'http://192.168.2.152:8253/IBAgent/exportFile'

    headers = {
        'Authorization': Data.Authorization,
        'Content-Type': 'application/json'  # 根据实际情况可能需要调整
    }

    data = {
        "fileId": Data.fileId,
        "_DownloadDataName": Data.DownloadDataName,
        "l_sign": Data.l_sign
    }

    # 将数据转换为 JSON 格式
    json_data = json.dumps(data)

    # 发送 POST 请求
    response = requests.post(url, headers=headers, data=json_data)
    # 检查响应状态码
    if response.status_code == 200:
        # 获取响应的字节流
        bytes_stream = response.content
        print(type(bytes_stream))
        # print(bytes_stream)

        # 将字节流转换为 BytesIO 对象
        bytes_io = io.BytesIO(bytes_stream)

    else:
        print(f"Failed to retrieve data, status code: {response.status_code}")


    # 使用 python-docx 库读取 Word 文档
    document = Document(bytes_io)

    # document = Document(Data.file_path)


    index_contract_dict = OrderedDict()
    # directory_index = float('inf')  # 设置为无限大
    directory_index = 0  # 设置为空
    for index, paragraph in enumerate(document.paragraphs):  # 获取目录后的索引段落字典 index_contract_dict = {段落索引:段落}
        if not paragraph.text.strip():
            continue
        if '第一节' in re.sub(' ', '', paragraph.text) and re.sub(' ', '', paragraph.text)[-1] not in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']:
            directory_index = index
            break
    for index, paragraph in enumerate(document.paragraphs):  # 获取目录后的索引段落字典 index_contract_dict = {段落索引:段落}
        if not paragraph.text.strip():
            continue
        if index >= directory_index:
            print(index, paragraph.text)
            # print('=========================================')
            index_contract_dict[int(index)] = paragraph.text.strip()

    for dict_index, (paragraph_index, contract) in enumerate(index_contract_dict.items()):  #
        text = contract.strip()
        pattern = r'[。！!？?]'
        # 使用re.split()函数切割字符串
        split_text = re.split(pattern, text)
        findall_text = re.findall(pattern, text)
        if split_text[-1] == '':
            split_text = split_text[:-1]
        for i in range(0, len(findall_text)):
            split_text[i] = split_text[i] + findall_text[i]
        # print(f'+++++++++++++++++++++{paragraph_index}+++++++++++++++++++++++++')
        # print(*split_text, sep='\n--------------------\n')
        for text_index, text in enumerate(split_text):
            # print(text,paragraph_index)
            await split_word_to_run(document.paragraphs[paragraph_index], text)

    new_file_stream = io.BytesIO()
    document.save(new_file_stream)
    new_file_stream.seek(0)
    # 指定保存文件的路径
    # 确保目录存在
    os.makedirs(file_path.parent, exist_ok=True)

    # 将内存中的流保存到本地文件
    with open(file_path, "wb") as file:
        file.write(new_file_stream.read())

    # 重置指针
    new_file_stream = io.BytesIO()
    document.save(new_file_stream)
    new_file_stream.seek(0)
    # 使用 getvalue() 方法获取字节流
    bytes_stream_again = new_file_stream.getvalue()

    e_time = time.time()
    print('转换时间', e_time - s_time)
    print(f"成功保存 {file_path}")

    return Response(content=bytes_stream_again, media_type="application/octet-stream")

class Data_C9(BaseModel):
    user_id: int = Field(default=0)
    file_name: str = Field(default="")
    bubble: dict = Field(default_factory=dict)
    fileId: str = Field(default="")
    status_code: int = Field(default=0)
    slice_index: int = Field(default=0)
    batch_slice_index: str = Field(default="")

@router.post("/C9_Insert_Bubble", summary="C9：输入：插入气泡的参数 => 插入气泡", include_in_schema=True)
async def C9_insert_bubble(data: Data_C9):
    db = LawDB()
    await db.init_conn()
    history_text = await db.insert_history(data.user_id, data.file_name, data.bubble, data.fileId, data.status_code, data.slice_index, data.batch_slice_index)
    return history_text


class Data_C10(BaseModel):
    user_id: int = Field(default=0)
    fileId: str = Field(default="")
    slice_index: int = Field(default=0)
    batch_slice_index: str = Field(default="")

@router.post("/C10_Delete_Bubble", summary="C10：输入：删除气泡的参数 => 删除气泡", include_in_schema=True)
async def C10_delete_bubble(data: Data_C10):
    db = LawDB()
    await db.init_conn()
    history_text = await db.delete_history(data.user_id, data.fileId, data.slice_index, data.batch_slice_index)
    return history_text


class Data_C11(BaseModel):
    user_id: int = Field(default=0)
    bubble: dict = Field(default_factory=dict)
    fileId: str = Field(default="")
    status_code: int = Field(default=0)
    slice_index: int = Field(default=0)
    batch_slice_index: str = Field(default="")

@router.post("/C11_Update_Bubble", summary="C11：输入：更新气泡的参数 => 更新气泡", include_in_schema=True)
async def C11_update_bubble(data: Data_C11):
    db = LawDB()
    await db.init_conn()
    history_text = await db.update_history(data.user_id, data.fileId, data.bubble, data.status_code, data.slice_index, data.batch_slice_index)
    return history_text


class Data_C12(BaseModel):
    user_id: int = Field(default=0)
    fileId: str = Field(default="")

@router.post("/C12_Select_Bubble", summary="C12：输入：查询气泡的参数 => 查询气泡", include_in_schema=True)
async def C12_select_bubble(data: Data_C12):
    db = LawDB()
    await db.init_conn()
    history_text = await db.select_history(data.user_id, data.fileId)
    return history_text


class Data_C13(BaseModel):
    content_text: str = Field(default="")
    material_list: List[dict] = Field(default_factory=list)

@router.post("/C13_Bubble_Score", summary="C13：输入：招股书切片，材料信息列表 => 评分值", include_in_schema=True)
async def C13_bubble_score(data: Data_C13):
    Global = Global_Score()
    score = await Global.bubble_score(data.content_text, data.material_list)
    return score

@router.post("/C14_PDF2JSON", summary="C14：输入pdf: => 输出: {段落索引: 段落文本}", include_in_schema=True)
async def C14_pdf2json(pdf: UploadFile = File(...)):

    try:
        # 检查文件扩展名是否为.pdf
        if pdf.filename.endswith(".pdf"):
            # 将文件保存到服务器上的临时文件中
            with open(f"./data/{pdf.filename}", "wb+") as file:
                shutil.copyfileobj(pdf.file, file)
            paragraph_dict = await get_full_text(str(pdf.filename)[:-4])
            return paragraph_dict
        else:
            return JSONResponse(content={"message": "File must be a PDF"}, status_code=400)
    except Exception as e:
        return JSONResponse(content={"message": str(e)}, status_code=500)


class Data_C15(BaseModel):
    content_text: str = Field(default="招股书切片")
    sector: str = Field(default="板块")

@router.post("/C15_Bubble_Title", summary="C15：输入: 招股书切片 => 输出: [问题分类标签]", include_in_schema=True)
async def C15_bubble_title(data: Data_C15):
    bubble = Bubble()
    problem_type_name_list = await bubble.title_bubble_v3(data.content_text, data.sector)
    return problem_type_name_list

class Data_C16(BaseModel):
    user_id: int = Field(default=0)
    file_name: str = Field(default="")
    bubble: dict = Field(default_factory=dict)
    fileId: str = Field(default="")
    status_code: int = Field(default=0)
    slice_index: int = Field(default=0)
    score: int = Field(default=0)

@router.post("/C16_Insert_Bubble_Overall", summary="C16：输入：插入全局气泡的参数 => 插入气泡", include_in_schema=True)
async def C16_insert_bubble_overall(data: Data_C16):
    db = LawDB()
    await db.init_conn()
    history_text = await db.insert_history_overall(data.user_id, data.file_name, data.bubble, data.fileId, data.status_code, data.slice_index, data.score)
    return history_text


class Data_C17(BaseModel):
    user_id: int = Field(default=0)
    fileId: str = Field(default="")
    slice_index: int = Field(default=0)

@router.post("/C17_Delete_Bubble_Overall", summary="C17：输入：删除全局气泡的参数 => 删除气泡", include_in_schema=True)
async def C17_delete_bubble_overall(data: Data_C17):
    db = LawDB()
    await db.init_conn()
    history_text = await db.delete_history_overall(data.user_id, data.fileId, data.slice_index)
    return history_text


class Data_C18(BaseModel):
    user_id: int = Field(default=0)
    bubble: dict = Field(default_factory=dict)
    fileId: str = Field(default="")
    status_code: int = Field(default=0)
    slice_index: int = Field(default=0)

@router.post("/C18_Update_Bubble_Overall", summary="C18：输入：更新全局气泡的参数 => 更新气泡", include_in_schema=True)
async def C18_update_bubble_overall(data: Data_C18):
    db = LawDB()
    await db.init_conn()
    history_text = await db.save_history_overall(data.user_id, data.fileId, data.bubble, data.status_code, data.slice_index)
    return history_text


class Data_C19(BaseModel):
    user_id: int = Field(default=0, description="用户id")
    fileId: str = Field(default="")

@router.post("/C19_Select_Bubble_Overall", summary="C19：输入：查询全局气泡的参数 => 查询气泡", include_in_schema=True)
async def C19_select_bubble_overall(data: Data_C19):
    db = LawDB()
    await db.init_conn()
    history_text = await db.select_history_overall(data.user_id, data.fileId)
    return history_text


class Data_C20(BaseModel):
    content_text: str = Field(default="")
    problem_id_list: List[int] = Field(default_factory=list)
    issuer_description: str = Field(default="")

@router.post("/C20_Get_Prompt_Third", summary="C20：输入: 招股书文本 和 问题ID列表 还有 C4->C5的输出（发行人说明/披露） => 输出: 提示词", include_in_schema=True)
async def C20_get_prompt_third_request(data: Data_C20):
    prompt = await get_prompt_third(data.content_text, data.problem_id_list, data.issuer_description)
    return prompt

class Data_C21(BaseModel):
    file_name: str = Field(default='', description="文件名")
    data: dict = Field(default_factory=dict, description='数据')

@router.post("/C21_Save_Mark", summary="C21：输入: 文件名 和 数据 => 存入标注数据", include_in_schema=True)
async def C21_save_mark(data: Data_C21):
    db = LawDB()
    await db.init_conn()
    save_text = await db.save_mark(data.file_name, data.data)
    return save_text

class Data_C22(BaseModel):
    file_name: str = Field(default='', description="文件名")

@router.post("/C22_Get_Mark", summary="C22：输入: 文件名 => 取出标注数据", include_in_schema=True)
async def C22_get_mark(data: Data_C22):
    db = LawDB()
    await db.init_conn()
    get_text = await db.get_mark(data.file_name)
    return get_text

class Data_C23(BaseModel):
    fileId: str = Field(default="", description="文件加密id")

@router.post("/C23_Select_History_Overall_Show", summary="C23：输入: 用户id 和 文件加密id => 查询所有需要显示的气泡", include_in_schema=True)
async def C23_select_history_overall_show(data: Data_C23):
    db = LawDB()
    await db.init_conn()
    history_text = await db.select_history_overall_show(data.fileId)
    return history_text

@router.post("/C29_Select_Task_Timeline", summary="C23：输入: 文件加密id => 获取任务总耗时", include_in_schema=True)
async def C29_select_task_timeline(data: Data_C23):
    db = LawDB()
    await db.init_conn()
    task_timeline = await db.select_task_timeline(data.fileId)
    return task_timeline

@router.post("/C30_Select_Rules_Reference", summary="C30: 文件加密id => 获取引用的规则列表")
async def C30_select_rules_reference(data: Data_C23):
    db = LawDB()
    await db.init_conn()
    rules_reference = await db.select_rules_reference(data.fileId)
    return rules_reference

@router.post("/C31_Select_File_Status", summary="C31: 文件加密id => 获取当前任务状态")
async def C31_select_file_status(data: Data_C23):
    """
    查询当前任务的状态
    :param data:
    :return:
    """
    db = db_pool.LawDB()
    await db.init_pool()
    task_status = await db.select_task(data.fileId)
    return IPOBaseResponse(status_code=200, message="任务状态查询成功", data=task_status, success=True)

class Data_C24(BaseModel):
    index_paragraph_sentence_dict: dict = Field(default_factory=dict, description='docx转json')
    sector: str = Field(default="", description="板块")
    industry: str = Field(default="", description="行业")
    region: str = Field(default="", description="地域")
    user_id: int = Field(default=0, description="用户id")
    file_name: str = Field(default="", description="招股书文件名.docx")
    fileId: str = Field(default="", description="文件加密id")

@router.post("/C24_Overall_Bubble", summary="C24：输入: docx转json 和 板块、行业、地域、用户id、招股书文件名、文件加密id => 输出: 全局气泡", include_in_schema=True)
async def C24_overall_bubble(data: Data_C24, background_tasks: BackgroundTasks):
    Overall = Overall_Bubble()
    background_tasks.add_task(Overall.overall_bubble, data.index_paragraph_sentence_dict, data.sector, data.industry, data.region, data.user_id, data.file_name, data.fileId)
    # 此处应该插入数据到task表，并且返回task_id
    return {"message": "全局请求已接收，后台任务正在执行"}

@router.post("/C24_Overall_Bubble_V2", summary="C24：输入: docx转json 和 板块、行业、地域、用户id、招股书文件名、文件加密id => 输出: 全局气泡", include_in_schema=True)
async def C24_overall_bubble(data: Data_C24, background_tasks: BackgroundTasks):
    Overall = Overall_Bubble_V2()
    background_tasks.add_task(Overall.overall_bubble, data.index_paragraph_sentence_dict, data.sector, data.industry, data.region, data.user_id, data.file_name, data.fileId)
    return {"message": "全局请求已接收，后台任务正在执行"}


class Data_C25(BaseModel):
    path: str = Field(default="", description="表格路径")

@router.post("/C25_Select_Table_Html", summary="C23：输入: 表格路径 => 表格Html", include_in_schema=True)
async def C25_select_table_html(data: Data_C25):
    db = LawDB()
    await db.init_conn()
    table_html = await db.select_table_html(data.path)
    return table_html

class Data_C26(BaseModel):
    title_dict: dict = Field(default_factory=dict, description="切片标题 与 问询标题 的字典")

@router.post("/C26_Slice_Query_Title", summary="C26：输入: {'切片标题': [], '问询标题' : [[], []]} => 筛选后的 {'切片标题': [], '问询标题' : [[], []]}", include_in_schema=True)
async def C26_slice_query_title(data: Data_C26):
    new_title_dict = await slice_query_title(data.title_dict)
    return new_title_dict


class Data_C27(BaseModel):
    index_paragraph_sentence_dict: dict = Field(default_factory=dict, description='招股书json')
    user_id: int = Field(default=0, description="用户id")
    fileId: str = Field(default="", description="文件加密id")

@router.post("/C27_Txt_Correction", summary="C27：输入: 招股书json => 文本纠错json", include_in_schema=True)
async def C27_text_correction(data: Data_C27, background_tasks: BackgroundTasks):
    Text = Text_Correction()
    background_tasks.add_task(Text.correction_text, data.index_paragraph_sentence_dict, data.user_id, data.fileId)
    return {"message": "文本纠错请求已接收，后台任务正在执行"}

class Data_C28(BaseModel):
    fileId: str = Field(default="", description="文件加密id")

@router.post("/C28_Select_Correction_Show", summary="C28：输入: 用户id 和 文件加密id => 查询所有需要显示的文本纠错气泡", include_in_schema=True)
async def C28_select_correction_show(data: Data_C28):
    db = LawDB()
    await db.init_conn()
    history_text = await db.select_correction_show(data.fileId)
    return history_text

class Data_C32(BaseModel):
    page: int = Field(default=1, description="页码")
    page_size: int = Field(default=20, description="每页数量")
    company_name: str = Field(default="", description="公司名称")
    company_abbreviation: str = Field(default="", description="公司简称")
    status: str = Field(default="", description="状态")
    stock_exchange: str = Field(default="", description="股票交易所")

# 展示所有问询的project，分页查询
@router.post("/C32_Select_Project_Show", summary="C32：输入: 页码、每页数量、公司名称、状态、交易所 => 查询筛选后的project", include_in_schema=True)
async def C32_select_project_show(data: Data_C32):
    db = LawDB()
    await db.init_conn()
    try:
        project_info = await db.select_project_show(
            page=data.page,
            page_size=data.page_size,
            company_name=data.company_name,
            company_abbreviation=data.company_abbreviation,
            status=data.status,
            stock_exchange=data.stock_exchange
        )
        return IPOBaseResponse(status_code=200, message="历史问询查询成功", data=project_info, success=True)
    finally:
        await db.close()  # 确保数据库连接关闭
